from unicodedata import category
import torch
import numpy as np

def default_collate(data_list):
    batch_size = len(data_list)
    images = [data['image'] for data in data_list]
    category_ids = [
        torch.FloatTensor([obj['category_id'] for obj in data['objects']])
        for data in data_list
    ]
    boxes = [
        torch.FloatTensor([obj['box'] for obj in data['objects']])
        for data in data_list
    ]
    images = torch.from_numpy(np.stack(images, axis=0))
    data = {
        'images': images,
        'category_ids': category_ids,
        'boxes': boxes,
    }
    return data
